import asyncio
import websockets
import json

async def connect_to_server():
    try:
        # 连接到 WebSocket 服务器
        async with websockets.connect('ws://localhost:8765') as websocket:
            print("WebSocket 连接已建立")
            
            # 发送 RTSP 地址给服务器
            video_path = 'rtsp://admin:zl888888@123.125.19.250:1554/Streaming/Channels/502'  # 请根据实际 RTSP 地址修改
            message = {"video_path": video_path}  # 将视频路径包装成 JSON 对象
            await websocket.send(json.dumps(message))  # 发送 JSON 格式的消息
            print(f"已发送视频路径: {video_path}")
            
            while True:
                try:
                    # 使用 asyncio.wait_for 设置超时
                    result = await asyncio.wait_for(websocket.recv(), timeout=30)  # 10秒超时
                except asyncio.TimeoutError:
                    print("接收数据超时")
                    break  # 可以选择退出或继续等待
                except websockets.exceptions.ConnectionClosed as e:
                    print(f"连接关闭: {e}")
                    break  # 退出循环
                
                # 解析 JSON 数据
                detections = json.loads(result)
                
                # 打印检测结果
                if "error" in detections:
                    print(f"错误: {detections['error']}")
                elif "message" in detections:
                    print(f"消息: {detections['message']}")
                else:
                    for detection in detections:
                        print(f"类别: {detection['class']}")
                        print(f"置信度: {detection['confidence']}")
                        print(f"坐标: {detection['coordinates']}")
                
                # 数据接收完毕，停止循环并关闭连接
                if not detections:
                    break  # 如果检测结果为空，退出循环
                await asyncio.sleep(1)  # 控制接收数据的速度
    except Exception as e:
        print(f"Error: {e}")
    finally:
        print("Closing WebSocket connection.")

# 启动客户端
async def main():
    await connect_to_server()

# 启动事件循环
if __name__ == "__main__":
    asyncio.run(main())
